import os
import json
import asyncio
from urllib.parse import quote
from dotenv import load_dotenv
from crawl4ai import AsyncWebCrawler, CrawlerRunConfig, LLMConfig
from crawl4ai.extraction_strategy import LLMExtractionStrategy
from pydantic import BaseModel, Field
from typing import List

# 1. 更新数据结构，增加'content'字段用于存储摘要或正文
class NewsArticle(BaseModel):
    title: str = Field(description="新闻的标题")
    url: str = Field(description="新闻的链接URL")
    content: str = Field(description="新闻的主要内容或摘要")

class NewsResult(BaseModel):
    articles: List[NewsArticle]


async def scrape_news(search_query: str, article_count: int) -> dict:
    """
    使用 Crawl4AI 和 LLM 提取策略来爬取新闻。
    :param search_query: 搜索关键词
    :param article_count: 要提取的新闻数量
    :return: 爬虫提取到的数据
    """
    print(f"为关键词 '{search_query}' 构造搜索URL...")
    encoded_query = quote(search_query)
    # 使用正确的客户端渲染URL
    search_url = f"https://so.news.cn/#search/0/{encoded_query}/1/0"
    #search_url = f"https://www.zhihu.com/search?type=content&q={encoded_query}"
    

    print("配置大语言模型 (LLM)...")
    llm_config = LLMConfig(
        provider="deepseek-chat",
        base_url="https://api.deepseek.com/v1",
        api_token=os.getenv("DEEPSEEK_API_KEY"),
    )

    print("配置LLM提取策略...")
    # 2. 优化指令，提供精确的HTML结构线索，并要求提取正文
    extraction_strategy = LLMExtractionStrategy(
        llm_config=llm_config,
        schema=NewsResult.model_json_schema(),
        instruction=(
            f"从以下HTML页面内容中，找到所有位于`<div class=\"result-item\">`标签内的新闻条目，并提取前 {article_count} 条。\n"
            f"对于每个新闻条目，请提取以下三项信息：\n"
            f"1. 标题：通常在`<h3 class=\"result-title\">`或类似的标题标签内。\n"
            f"2. 链接：从标题的`<a>`标签中获取`href`属性。\n"
            f"3. 正文：从`<p class=\"result-abstract\">`标签中提取新闻的摘要或主要描述文字。\n"
            f"请将所有提取到的新闻整理成一个列表返回。"
        ),
        response_format={"type": "json_object"}
    )

    # 3. 配置爬虫精准等待，直到新闻列表加载完成
    run_config = CrawlerRunConfig(
        extraction_strategy=extraction_strategy,
        # 强制爬虫等待，直到第一个 class="result-item" 的元素出现
        wait_for="div.result-item"
    )

    print("初始化并运行 Playwright 爬虫...")
    async with AsyncWebCrawler() as crawler:
        result = await crawler.arun(
            url=search_url,
            config=run_config,
        )
        
        print("爬取任务完成。")
        if result.extracted_content:
            try:
                return json.loads(result.extracted_content)
            except json.JSONDecodeError:
                print("LLM返回的不是有效的JSON格式。")
                print("原始返回内容:", result.extracted_content)
                return None
        return None

async def main():
    """
    主函数，定义爬取目标并处理结果。
    """
    SEARCH_QUERY = "亦庄"
    ARTICLE_COUNT = 5

    try:
        extracted_data = await scrape_news(SEARCH_QUERY, ARTICLE_COUNT)

        if not extracted_data:
            print("\n未能提取到任何数据，请检查爬虫日志或目标网页。")
            return

        articles_list = []
        if isinstance(extracted_data, dict) and 'articles' in extracted_data:
            articles_list = extracted_data['articles']
        elif isinstance(extracted_data, list):
            articles_list = extracted_data
        
        if not articles_list:
            print("\n提取到的数据为空或格式不正确。")
            return

        output_dir = "output"
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, "news_data.txt")

        with open(output_path, "w", encoding="utf-8") as f:
            print(f"\n正在将数据写入到: {output_path}")
            # 4. 调整文件写入逻辑，增加正文内容
            for i, article in enumerate(articles_list):
                title = article.get('title', '无标题')
                url = article.get('url', '无链接')
                content = article.get('content', '无正文')
                
                f.write(f"新闻 {i+1}:\n")
                f.write(f"  标题: {title}\n")
                f.write(f"  链接: {url}\n")
                f.write(f"  正文: {content}\n")
                f.write("-" * 20 + "\n")

        print(f"\n成功! 爬取的数据已保存为TXT文件: {output_path}")
        with open(output_path, 'r', encoding='utf-8') as f:
            print("文件内容预览:")
            print(f.read())

    except Exception as e:
        print(f"\n执行过程中发生错误: {e}")

if __name__ == "__main__":
    load_dotenv()
    api_key = os.getenv("DEEPSEEK_API_KEY")
    if not api_key:
        raise ValueError(
            "未找到 Deepseek API Key。请确保您的项目中有一个 .env 文件，"
            "并且其中包含 'DEEPSEEK_API_KEY=your_key'。"
        )
    asyncio.run(main())